tabmlp_keep_keys = [
    "mlp_hidden_dims",
    "mlp_activation",
    "mlp_dropout",
    "mlp_batchnorm",
    "mlp_batchnorm_last",
    "mlp_linear_first",
    "embed_dropout",
    "lr",
    "batch_size",
    "weight_decay",
    "optimizer",
    "lr_scheduler",
    "base_lr",
    "max_lr",
    "div_factor",
    "final_div_factor",
    "n_cycles",
]


tabresnet_keep_keys = [
    "blocks_dims",
    "blocks_dropout",
    "mlp_hidden_dims",
    "mlp_activation",
    "mlp_dropout",
    "mlp_batchnorm",
    "mlp_batchnorm_last",
    "mlp_linear_first",
    "embed_dropout",
    "lr",
    "batch_size",
    "weight_decay",
    "optimizer",
    "lr_scheduler",
    "base_lr",
    "max_lr",
    "div_factor",
    "final_div_factor",
    "n_cycles",
]


tabnet_keep_keys = [
    "n_steps",
    "step_dim",
    "attn_dim",
    "ghost_bn",
    "virtual_batch_size",
    "momentum",
    "gamma",
    "dropout",
    "embed_dropout",
    "lr",
    "batch_size",
    "weight_decay",
    "lambda_sparse",
    "optimizer",
    "lr_scheduler",
    "base_lr",
    "max_lr",
    "div_factor",
    "final_div_factor",
    "n_cycles",
]


tabtransformer_keep_keys = [
    "embed_dropout",
    "full_embed_dropout",
    "shared_embed",
    "add_shared_embed",
    "frac_shared_embed",
    "input_dim",
    "n_heads",
    "n_blocks",
    "dropout",
    "ff_hidden_dim",
    "transformer_activation",
    "mlp_hidden_dims",
    "mlp_activation",
    "mlp_batchnorm",
    "mlp_batchnorm_last",
    "mlp_linear_first",
    "with_wide",
    "lr",
    "batch_size",
    "weight_decay",
    "optimizer",
    "lr_scheduler",
    "base_lr",
    "max_lr",
    "div_factor",
    "final_div_factor",
    "n_cycles",
]
